import os
import time

from flask import Flask, request, flash, url_for, send_from_directory
from werkzeug.utils import redirect, secure_filename
import requests
import json
import cv2
import base64
import uuid

app = Flask(__name__)

ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png'}


def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def cv2_to_base64(image):
    data = cv2.imencode('.jpg', image)[1]
    return base64.b64encode(data.tostring()).decode('utf8')


@app.route('/download/<filename>', methods=['GET'])
def download(filename):
    return send_from_directory("D://style_transfer2//", filename, as_attachment=True)


@app.route('/getOcr/<filename>', methods=['GET'])
def getOcr(filename):
    data = {'images': [cv2_to_base64(cv2.imread("D://style_transfer2//" + filename))]}
    headers = {"Content-type": "application/json"}
    url = "http://127.0.0.1:8866/predict/chinese_ocr_db_crnn_server"
    r = requests.post(url=url, headers=headers, data=json.dumps(data))
    return r.json()


@app.route('/upload', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        # check if the post request has the file part
        if 'file' not in request.files:
            flash('No file part')
            return redirect(request.url)
        file = request.files['file']
        # If the user does not select a file, the browser submits an
        # empty file without a filename.
        if file.filename == '':
            flash('No selected file')
            return redirect(request.url)
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            file_type_index = filename.rfind('.')
            file_type = filename[file_type_index:]
            new_file_name = time.time().__str__().replace('.', '') + file_type
            file.save(os.path.join('D:/style_transfer2', new_file_name))
            return new_file_name
        return '文件格式不允许'


if __name__ == '__main__':
    app.run(debug=True, port=8001)
